import logging
import os

import docker
import git
import yaml
from docker.errors import ImageNotFound, NotFound

from app.core.flask_celery import ContextTask
from app.extensions import celery, db
from app.models.admin import Config
from app.models.docker import DockerResource, DockerRunner
from app.services.docker import get_free_port
from app.services.system import create_admin_message
from app.utils.tools import find_directories_with_filename

logger = logging.getLogger(__name__)


@celery.task(base=ContextTask)
def stop_vulnerability_resource(run_id):
    """
    销毁漏洞环境
    """
    logger.info("start destroy vulnerability runner {}".format(run_id))
    runner: DockerRunner = DockerRunner.get_by_id(run_id)
    docker_api = Config.get_config(Config.KEY_DOCKER_API)
    if not runner:
        return
    client = docker.DockerClient(docker_api)
    try:
        docker_container = client.containers.get(runner.container_id)
        docker_container.stop()
        docker_container.remove()
    except docker.errors.NotFound as e:
        logger.error(e)
    db.session.delete(runner)
    db.session.commit()
    logger.info("destroy vulnerability runner:{}".format(run_id))


def start_vuln_resource(resource_id, user_id=None, admin_id=None):
    """
    启动docker 资源
    """
    if user_id is None and admin_id is None:
        raise ValueError()
    resource = DockerResource.get_by_id(resource_id)
    client = docker.DockerClient(Config.get_config(Config.KEY_DOCKER_API))
    try:
        image = client.images.get(resource.image)
    except ImageNotFound as e:
        logger.exception(e)
        raise ValueError("当前题目环境缺失、请联系管理员！")
    # 解析镜像端口
    image_config = image.attrs["Config"]
    port_range = Config.get_config(Config.KEY_PORT_RANGE)
    try:
        start_port, end_port = port_range.split("-")
        start_port, end_port = int(start_port), int(end_port)
        assert start_port < end_port
    except (AssertionError, ValueError, IndexError):
        raise ValueError("服务器缺少资源、请联系管理员")
    port_dict = {}
    if resource.ports:
        for docker_port in resource.ports.split(","):
            port_dict[docker_port] = get_free_port(start_port, end_port)
    elif "ExposedPorts" in image_config:
        port_dict = image_config["ExposedPorts"]
        for docker_port, host_port in port_dict.items():
            port_dict[docker_port] = get_free_port(start_port, end_port)
    image_name = image.attrs["RepoTags"][0].replace(":", ".")
    container_name = f"{image_name}_{user_id}".replace("/", "-")
    # 检查docker 是否已存在
    try:
        c = client.containers.get(container_name)
        c.stop()
        c.remove()
    except docker.errors.NotFound:
        pass
    try:
        docker_container = client.containers.run(
            image=image.id, name=container_name, ports=port_dict, detach=True
        )
    except docker.errors.APIError as e:
        logger.exception(e)
        raise ValueError("题目启动失败")
    # 查看是否有历史记录
    try:
        docker_runner = (
            db.session.query(DockerRunner)
            .filter(DockerRunner.name == docker_container.name)
            .first()
        )
        if docker_runner:
            docker_runner.container_id = docker_container.id
            docker_runner.port_info = port_dict
            docker_runner.save()
        else:
            docker_runner = DockerRunner.create(
                name=docker_container.name,
                resource_id=resource_id,
                admin_id=admin_id,
                user_id=user_id,
                port_info=port_dict,
                container_id=docker_container.id,
            )
        return docker_runner.id
    except Exception as e:
        logger.exception(e)
        db.session.rollback()


@celery.task(base=ContextTask)
def sync_remote_vulnerability_repo(repo, admin_id=None, max_retry=3):
    logger.info(f"start sync:{repo}")
    repo_name = repo.split("/")[-1].split(".")[0]
    username = repo.split("/")[-2]
    local_repo = f"/opt/vulnerability/{username}/{repo_name}"
    # 判断本地目录是否存在
    try:
        if os.path.exists(local_repo):
            logger.info(f"Pull To {local_repo}")
            git.Repo(local_repo).git.pull()
        else:
            logger.info(f"Clone {repo} to {local_repo}")
            git.Repo.clone_from(repo, local_repo)
        logger.info("Pull Successful")
    except git.exc.GitCommandError as e:
        logger.error(e)
        # retry
        if max_retry <= 0:
            logger.error("同步仓库失败")
            create_admin_message(admin_id, f"同步远程漏洞仓库失败\n{e}")
            return
        sync_remote_vulnerability_repo.delay(repo, admin_id, max_retry - 1)
        return
    # scan directory
    client = docker.DockerClient(Config.get_config(Config.KEY_DOCKER_API))
    vulnerabilities = find_directories_with_filename(local_repo, filename="metadata.yml")
    for directory in vulnerabilities:
        logger.info("update {} ...".format(directory))
        with open(os.path.join(directory, "metadata.yml")) as f:
            yaml_data = yaml.safe_load(f)
        image = yaml_data["image"]
        description = yaml_data["description"]
        name = yaml_data["name"]
        ports = ",".join([str(i) for i in yaml_data.get("ports", [])])
        if description.startswith("file:"):
            desc_file = os.path.join(directory, description.lstrip("file:"))
            if os.path.exists(desc_file):
                with open(desc_file, "r", encoding="UTF-8") as f:
                    description = f.read()
        # build
        obj = db.session.query(DockerResource).filter(DockerResource.image == image).first()
        if obj:
            logger.info(f"Image:{image} Is Already")
            obj.description = description
            obj.ports = ports
            db.session.commit()
            continue
        else:
            obj = db.session.query(DockerResource).filter(DockerResource.name == name).first()
            if obj:
                obj.ports = ports
                obj.image = image
            else:
                obj = DockerResource(
                    name=name,
                    resource_type="VUL",
                    image=image,
                    description=description,
                    ports=ports,
                    cve=yaml_data.get("cve", []),
                    app=yaml_data.get("app"),
                )
        docker_file = os.path.join(directory, "Dockerfile")
        if os.path.exists(docker_file):
            # build
            obj.status = DockerResource.STATUS_BUILD
            obj.dockerfile = docker_file
            obj.docker_type = DockerResource.DOCKER_TYPE_LOCAL_IMAGE
            # 判断镜像是否存在 如果存在就标记状态 已就绪
            try:
                client.images.get(image)
                obj.status = DockerResource.STATUS_BUILD
            except NotFound:
                obj.status = DockerResource.STATUS_INIT
        else:
            obj.description = description
            obj.status = DockerResource.STATUS_INIT
            obj.docker_type = DockerResource.DOCKER_TYPE_REMOTE_IMAGE
        db.session.add(obj)
        db.session.commit()
        logger.info("Add Image:{}".format(image))
    create_admin_message(admin_id, "同步远程漏洞仓库完成")
